# 预测提交
import os
import torch
import numpy as np
import pandas as pd
from tqdm import tqdm
from public import (
    DataConfig, ModelConfig, RNAModel, 
    read_fasta_biopython, featurize
)

# 配置参数
class RunConfig:
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model_path = "RNA/RNA_design_public/ckpts/public_v2/best_epoch70_lr0.0001.pt"
    data_config = DataConfig(
        test_npy_data_dir="../saisdata/coords",
        test_data_path=""  # 实际从文件读取，不需要csv
    )

def load_test_data():
    """加载测试数据"""
    coords_dir = RunConfig.data_config.test_npy_data_dir
    seqs_dir = "../saisdata/seqs"
    
    samples = []
    for fname in os.listdir(coords_dir):
        if fname.endswith(".npy"):
            pdb_id = os.path.splitext(fname)[0]
            # 加载坐标数据
            coords = np.load(os.path.join(coords_dir, fname))
            # 加载假序列（仅用于占位）
            fasta_path = os.path.join(seqs_dir, f"{pdb_id}.fasta")
            seq = list(read_fasta_biopython(fasta_path).values())[0]
            samples.append({
                "name": pdb_id,
                "coords": {"P": coords[:,0], "O5'": coords[:,1], 
                          "C5'": coords[:,2], "C4'": coords[:,3],
                          "C3'": coords[:,4], "O3'": coords[:,5]},
                "seq": seq
            })
    return samples

def predict():
    # 初始化模型
    model = RNAModel(ModelConfig()).to(RunConfig.device)
    model.load_state_dict(torch.load(RunConfig.model_path))
    model.eval()

    # 加载数据
    test_samples = load_test_data()
    test_loader = DataLoader(
        test_samples, 
        batch_size=16,
        collate_fn=featurize,
        num_workers=0
    )

    # 执行预测
    results = []
    with torch.no_grad():
        for batch in tqdm(test_loader, desc="Predicting"):
            X, S, mask, lengths, names = batch
            X = X.to(RunConfig.device)
            mask = mask.to(RunConfig.device)
            
            # 使用虚拟序列占位
            S_dummy = torch.zeros_like(S).to(RunConfig.device)
            
            # 获取预测结果
            logits, _, _ = model(X, S_dummy, mask)
            preds = torch.argmax(logits, dim=-1).cpu().numpy()

            # 转换核苷酸序列
            start_idx = 0
            alphabet = 'AUCG'
            for i, length in enumerate(lengths):
                end_idx = start_idx + length
                seq = ''.join([alphabet[p] for p in preds[start_idx:end_idx]])
                results.append({"pdb_id": names[i], "seq": seq})
                start_idx = end_idx

    # 保存结果
    df = pd.DataFrame(results)
    os.makedirs("../saisresult", exist_ok=True)
    df.to_csv("../saisresult/submit.csv", index=False)
    print(f"预测完成！结果已保存至 ../saisresult/submit.csv")

if __name__ == "__main__":
    predict()